import rospy
from gazebo_msgs.srv import SetModelState
from gazebo_msgs.msg import ModelState
import math
import threading

from enu_xyz import GNC_COORDINATE
from queue import Queue
import numpy as np
import ROSinfo
import torch
import time

from DDPG_ import PolicyNetwork
import matplotlib.pyplot as plt
# 在全局范围内创建一个队列
boat_pos_queue = Queue()
boat_yaw_queue = Queue()
boat_vel_queue = Queue()
boat_speed_queue = Queue()
M_PI = math.pi/3
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

# 输出船的舵角和推力

class LADRC_NUM:
    def __init__(self, r=2.0, w0=0.0, wc=0.0, b0=0.0):
        self.v1 = 0
        self.v2 = 0
        self.v3 = 0
        self.r = r
        self.h = 0.001
        self.z1 = 0
        self.z2 = 0
        self.z3 = 0
        self.w0 = w0
        self.wc = wc
        self.b0 = b0
        self.u = 0

class LADRC:
    def __init__(self, ladrc_num):
        self.ladrc_num = ladrc_num


    def LADRC_TD(self, expect):
        fh = -1 * self.ladrc_num.r * self.ladrc_num.r * (self.ladrc_num.v1 - expect) - 2 * self.ladrc_num.r * self.ladrc_num.v2
        self.ladrc_num.v1 += self.ladrc_num.v2 * self.ladrc_num.h
        self.ladrc_num.v2 += fh * self.ladrc_num.h
        # self.ladrc_num.v1 = expect

    def LADRC_ESO(self, feedback):
        Beita_01 = 3 * self.ladrc_num.w0
        # print(self.ladrc_num.w0)
        Beita_02 = 3 * self.ladrc_num.w0 * self.ladrc_num.w0
        Beita_03 = self.ladrc_num.w0 * self.ladrc_num.w0 * self.ladrc_num.w0

        e = self.ladrc_num.z1 - feedback
        self.ladrc_num.z1 += (self.ladrc_num.z2 - Beita_01 * e) * self.ladrc_num.h
        self.ladrc_num.z2 += (self.ladrc_num.z3 - Beita_02 * e + self.ladrc_num.b0 * self.ladrc_num.u) * self.ladrc_num.h
        self.ladrc_num.z3 += -1 * Beita_03 * e * self.ladrc_num.h

    def LADRC_LF(self):
        Kp = self.ladrc_num.wc * self.ladrc_num.wc
        Kd = 2 * self.ladrc_num.wc
        e1 = self.ladrc_num.v1 - self.ladrc_num.z1
        e2 = -1 * self.ladrc_num.z2
        u0 = Kp * e1 + Kd * e2
        self.ladrc_num.u = (u0 - self.ladrc_num.z3) / self.ladrc_num.b0

    def LADRC_Loop(self, expect_value, measure):
        self.LADRC_TD(expect_value)
        self.LADRC_ESO(measure)
        self.LADRC_LF()


def plot(frame_idx, rewards):
    plt.figure(figsize=(5, 5))
    # plt.subplot(131)
    plt.title('frame %s. reward: %s' % (frame_idx, rewards[-1]))
    plt.plot(rewards)
    plt.show()

def calculate_metrics(yaw, origin_x, origin_y, hangdian2_x, hangdian2_y, boat_x, boat_y):
    yaw_rad = math.radians(yaw)

    # 计算船在指定向量方向上的速度投影
    # vector_hangdian = np.array([hangdian2_x - origin_x, hangdian2_y - origin_y])

    # v_x = np.array([forward_v * math.cos(yaw_rad), forward_v * math.sin(yaw_rad)])
    # v_y = np.array([lateral_v * math.cos(yaw_rad), lateral_v * math.sin(yaw_rad)])
    # velocity_projection = (np.dot(v_x, vector_hangdian) / math.sqrt(np.dot(vector_hangdian, vector_hangdian)) )
    #                        # + np.dot(v_y, vector_hangdian) / math.sqrt(np.dot(vector_hangdian, vector_hangdian)))

    # 计算船到指定直线的距离
    line_vector_x = hangdian2_x - origin_x
    line_vector_y = hangdian2_y - origin_y
    boat_vector_x = boat_x - origin_x
    boat_vector_y = boat_y - origin_y
    distance_to_line = (line_vector_x * boat_vector_y - line_vector_y * boat_vector_x) / math.sqrt(line_vector_x**2 + line_vector_y**2)
    # 计算矢量场制导律的角度
    k = 0.1
    ang_ = -math.atan(k * distance_to_line) + math.atan2(line_vector_y, line_vector_x)
    # 计算ang_和yaw_rad的夹角，限制在-pi到pi之间
    angle_error = ang_ - yaw_rad

    # 将角度误差调整到 -π 到 π 之间
    angle_error = (angle_error + math.pi) % (2 * math.pi) - math.pi

    return distance_to_line, angle_error

def calculate_circular_metrics(yaw, circle_x, circle_y, radius, boat_x, boat_y):
    yaw_rad = math.radians(yaw)

    # 计算船到圆心的距离
    boat_to_center_x = boat_x - circle_x
    boat_to_center_y = boat_y - circle_y
    distance_to_center = math.sqrt(boat_to_center_x**2 + boat_to_center_y**2)

    # 计算船到圆周的误差距离（正值表示在外部，负值表示在内部）
    distance_to_circle = distance_to_center - radius

    # 计算船相对圆周的切线方向（目标航向）
    tangent_angle = math.atan2(boat_to_center_y, boat_to_center_x) - math.pi / 2  # 顺时针方向
    # 若希望逆时针绕圈，可改为 -math.pi / 2
    print(tangent_angle * 180 /math.pi)
    # 计算矢量场制导律的角度
    k = 0.12  # 导数增益系数，可调整
    ang_ = tangent_angle - math.atan(k * distance_to_circle)

    # 计算角度误差，限制在 -π 到 π 之间
    angle_error = ang_ - yaw_rad
    angle_error = (angle_error + math.pi) % (2 * math.pi) - math.pi
    print(angle_error * 180 / math.pi,boat_x,boat_y)
    return distance_to_circle, angle_error

def random_target():
    target_x = 0
    target_y = 0
    while (target_x ** 2 + target_y ** 2) < 36:
        target_x = np.random.uniform(-10, 10)
        target_y = np.random.uniform(-10, 10)
    return target_x, target_y

def goto_point(given_point):
    # 等待服务可用
    rospy.wait_for_service('/gazebo/set_model_state')
    try:
        # 创建服务代理
        set_state = rospy.ServiceProxy('/gazebo/set_model_state', SetModelState)

        # 构造状态消息
        state_msg = ModelState()
        state_msg.model_name = 'wamv'
        state_msg.pose.position.x = given_point[0]
        state_msg.pose.position.y = given_point[1]
        state_msg.pose.position.z = 0.0  # 假设 z 坐标为 0
        # 保持默认朝向（或者根据需要设置四元数）
        theta = math.radians(150)  # 150度 = 5π/6 弧度
        state_msg.pose.orientation.x = 0.0
        state_msg.pose.orientation.y = 0.0
        state_msg.pose.orientation.z = math.sin(theta / 2)
        state_msg.pose.orientation.w = math.cos(theta / 2)
        state_msg.reference_frame = 'world'

        # 调用服务进行状态设置（瞬移）
        resp = set_state(state_msg)
        print("无人船已瞬移到目标点：", given_point)
    except rospy.ServiceException as e:
        print("调用 /gazebo/set_model_state 服务失败: %s" % e)

def plot_boat_coordinates(boat_coordinates, string):
    # 提取 x 和 y 坐标
    x_coordinates, y_coordinates = zip(*boat_coordinates)

    # 绘制船的轨迹
    plt.plot(x_coordinates, y_coordinates, marker='o', linestyle='-')
    plt.axis('auto')  # 自动调整坐标轴范围
    # 设置图形属性
    plt.title(string)
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)

    # 显示图形
    plt.show()
def map_action_to_parameters(action):
    """
    将动作向量映射到ADRC参数的范围内。

    参数:
        action (np.array 或 list): 长度为2，取值范围[-1, 1]

    返回:
        tuple: (w0, b0)
    """

    # # w0 在 [100, 300] 之间浮动
    # w0 = (action[0] + 1) / 2 * (400 - 100) + 100
    # # b0 在 [750, 2250] 之间浮动
    # b0 = (action[1] + 1) / 2 * (2250 - 750) + 750

    # w0 = (action[0] + 1) / 2 * (450 - 50) + 50
    # # b0 在 [750, 2250] 之间浮动
    # b0 = (action[1] + 1) / 2 * (3000 - 750) + 750
    w0 = (action[0] + 1) / 2 * (450 - 50) + 50
    # b0 在 [750, 2250] 之间浮动
    b0 = (action[1] + 1) / 2 * (2500 - 750) + 750

    return w0, b0


if __name__ == "__main__":

    gnc_coordinate = GNC_COORDINATE()

    # 训练路径航点
    # hangdian1_lon = 150.6758493
    # hangdian1_lat = -33.7219811
    hangdian1_lon = 150.6769266
    hangdian1_lat = -33.7228700
    hangdian2_lon = 150.6769266
    hangdian2_lat = -33.7228700
    # 设置初始点经纬度
    origin_lon = hangdian1_lon
    origin_lat = hangdian1_lat
    origin_x = gnc_coordinate.Lon_trans_to_E(hangdian1_lon, hangdian1_lat , origin_lon, origin_lat)
    origin_y = gnc_coordinate.Lat_trans_to_N(hangdian1_lon, hangdian1_lat , origin_lon, origin_lat)
    hangdian2_x = gnc_coordinate.Lon_trans_to_E(hangdian2_lon, hangdian2_lat , origin_lon, origin_lat)
    hangdian2_y = gnc_coordinate.Lat_trans_to_N(hangdian2_lon, hangdian2_lat , origin_lon, origin_lat)

    def update_func_gps(handle, lat, lon, alt):
        # 将ROS传回的经纬度转换为坐标系
        boat_x = gnc_coordinate.Lon_trans_to_E(lon, lat, origin_lon, origin_lat)
        boat_y = gnc_coordinate.Lat_trans_to_N(lon, lat, origin_lon, origin_lat)
        boat_pos_queue.put((boat_x, boat_y))


    def update_func_ea(handle, yaw, pitch, roll):
        boat_yaw_queue.put(yaw)

    def update_func_ground_speed(handle, speed , forward_v , lateral_v):
        boat_vel_queue.put((forward_v,lateral_v))
        boat_speed_queue.put(speed)

    def update_func_gyro(handle, gyroX, gyroY, gyroZ):
        pass

    def update_func_acc(handle, accX, accY, accZ):
        pass

    rospy.init_node('my_ROSVRX_node')

    update_handle = None

    ros_vrx_info = ROSinfo.ROSVRXInfo(update_func_gps, update_func_ea, update_func_ground_speed,
                              update_func_gyro, update_func_acc, update_handle)

    propeller_mid = 0
    servo_mid = 0
    rosvrx_control = ROSinfo.propulsion_VRX(propeller_mid, servo_mid)

    # 在主函数中调用 manual_act 控制无人船

    def publish_commands():
        import os
        # if not os.path.exists("4_load_data.txt"):
        #     with open("data.txt", "w") as f:
        #         f.write("三级海况下TD3跟踪\n")

        while not rospy.is_shutdown():
            boat_coordinates = []
            dis_error =[]
            angs = []
            angs_error = []
            power_left = []
            power_right = []
            servo = []
            # 下载模型参数 并放到模型中
            state_dim = 4  # 根据实际情况设置状态维度
            action_dim = 2  # 根据实际情况设置动作维度
            hidden_dim = 128  # 根据实际情况设置隐藏层维度
            loaded_model = PolicyNetwork(state_dim, action_dim, hidden_dim)
            loaded_model.load_state_dict(torch.load('/home/ltp/train model/5/三级海况800轮加跟踪误差.pt'))
            # 将模型设置为评估模式
            loaded_model.eval()
            # 提取船当前的状态
            ADRC_param = LADRC_NUM(r=2, w0=200, wc=50, b0=1500)
            ADRC_ctl = LADRC(ADRC_param)
            goto_point([-260,190])   #  goto_point([-260,150])
            time.sleep(1)
            boat_x, boat_y = boat_pos_queue.get()
            print(boat_x, boat_y)
            first_time =time.time()
            # 直线
            # hd_list = [(0,0),(50,80),(180,100),(280,60),(280,0)]
            hd_list = [(-10, 50)]

            for i, coordinates in enumerate(hd_list):
                x, y = coordinates
                if i > 0 :
                    prev_x, prev_y = hd_list[i - 1]
                else:
                    prev_x, prev_y = origin_x, origin_y

                pre_ang_error = 0
                pre_t_error = 0

                # time.sleep(10)
                pre_time = time.time()
                # 直线
                # while (x- boat_x)**2+(y - boat_y)**2 > 36:
                for step in range(3000):
                    boat_x, boat_y = boat_pos_queue.get()
                    boat_coordinates.append((boat_x, boat_y))
                    yaw = boat_yaw_queue.get()
                    speed = boat_speed_queue.get()
                    # 直线
                    # t_error, ang_error = calculate_metrics(yaw, prev_x, prev_y,
                    #                                        x, y, boat_x, boat_y)
                    # 圆
                    t_error, ang_error = calculate_circular_metrics(yaw, 50, 50,
                                                            60, boat_x, boat_y)
                    angle_error_deriv = abs(ang_error) - abs(pre_ang_error)
                    pre_ang_error = ang_error
                    # input = torch.from_numpy(np.array([t_error, ang_error, dt_t_error, dt_ang_error] + action)).double()
                    input = np.array([t_error, ang_error, angle_error_deriv, ADRC_param.z3])
                    input = torch.FloatTensor(input).unsqueeze(0).to(device)
                    with torch.no_grad():
                        output_action = loaded_model(input)
                    output_action = np.squeeze(output_action.numpy())
                    w0, b0 = map_action_to_parameters(output_action)
                    ADRC_param.w0 = 450
                    ADRC_param.wc = 450 / 4  # wc = w0/4 的逻辑
                    ADRC_param.b0 = 1500
                    ADRC_ctl.LADRC_Loop(0, ang_error)
                    turn = ADRC_param.u

                    rosvrx_control.manual_act(0.8, 0.8, turn)
                    with open("4_load_data.txt", "a") as f:
                        # 这里的 angle_error = next_state[0], angle_error_deriv = next_state[1]
                        f.write(f"速度: {speed}, 路径误差: {t_error}, 角度误差(度): {ang_error * 180 / math.pi}, 当前位置: {boat_x},{boat_y}, w0: {w0}, b0: {b0}, z3:{ADRC_param.z3}\n")
                    current_time = time.time()
                    dis_error.append((current_time - first_time,t_error))
                    angs.append(((current_time - first_time),yaw))
                    angs_error.append(((current_time - first_time),ang_error * 180 / math.pi))
            # 从 angs_error 中提取 ang_error 列表
            ang_errors = [ang_error for _, ang_error in angs_error]
            # 使用 NumPy 计算方差
            variance = np.var(ang_errors)
            print("方差为:", variance)
            plot_boat_coordinates(boat_coordinates,'无人船位置图')
            plot_boat_coordinates(dis_error,'无人船误差图')
            plot_boat_coordinates(angs_error,'无人船角度误差图')
    # 创建一个线程来发布命令
    publish_thread = threading.Thread(target=publish_commands)
    publish_thread.start()

    try:
        ros_vrx_info.ROSVRX_info_create()

    except rospy.ROSInterruptException:
        pass
